<#
.SYNOPSIS
    Configure Logon/Logoff Audit Policies with UI Visibility - Updates both auditpol and audit.csv database.

.SCRIPTTYPE
    Computer Configuration

.DESCRIPTION
    This script applies Logon/Logoff audit policy settings by:
    1. Setting active audit policy via auditpol.exe (makes policies active)
    2. Updating Group Policy audit.csv database (makes policies visible in secpol.msc UI)
    
    The script supports simplified audit policy configuration using comprehensive logging, validation, and error handling.
    
    INPUT FORMAT:
        "1"     → Success audit only (logs successful events)
        "0"     → Failure audit only (logs failed events)  
        "1,0"   → Both success and failure auditing (logs all events)
        ""      → No auditing (skip this policy)
    
    AUDIT POLICIES (11 policies in order):
        1.  Audit Logon - Tracks user logon events and authentication
        2.  Audit Logoff - Monitors user logoff and session termination  
        3.  Audit Account Lockout - Captures account lockout events
        4.  Audit IPsec Main Mode - Tracks IPsec main mode negotiations
        5.  Audit IPsec Quick Mode - Monitors IPsec quick mode operations
        6.  Audit IPsec Extended Mode - Captures IPsec extended mode events
        7.  Audit Special Logon - Tracks special logon types and service accounts
        8.  Audit Other Logon/Logoff Events - Monitors miscellaneous logon/logoff events
        9.  Audit Network Policy Server - Captures NPS/RADIUS authentication events
        10. Audit User / Device Claims - Tracks user and device claims in authentication
        11. Audit Group Membership - Monitors group membership in tokens

    SECURITY RECOMMENDATIONS:
        - Logon: Enable both for comprehensive logon monitoring
        - Logoff: Enable success for session tracking
        - Account Lockout: Enable both for lockout attack detection
        - IPsec policies: Enable failure for IPsec troubleshooting
        - Special Logon: Enable both for privileged account monitoring
        - Other Logon/Logoff Events: Enable both for complete visibility
        - Network Policy Server: Enable both for RADIUS monitoring
        - User/Device Claims: Enable success for claims-based auth tracking
        - Group Membership: Enable success for token analysis

.PARAMETER PolicyValues
    JSON array string containing 11 audit values (in order). Use "1" for success, "0" for failure, "1,0" for both, or "" to skip.
    Format: '["value1","value2",...,"value11"]'

.PARAMETER LogLevel
    Set logging verbosity: Silent, Normal, Verbose, Debug

.PARAMETER LogPath
    Path for log file output. If not specified, auto-configures to agent logs directory.

.PARAMETER WhatIf
    Preview changes without applying them.

.EXAMPLE
    .\Set-LogonLogoffAuditPolicies.ps1 '["1,0","1","1,0","","","","1","","1,0","",""]'
    Configures comprehensive logon monitoring with selective policies
        
.EXAMPLE
    .\Set-LogonLogoffAuditPolicies.ps1 '["1","1","1","0","0","0","1","1","1","1","1"]' -WhatIf
    Preview mode for all 11 policies

.NOTES
    - Requires administrative privileges (Run as Administrator)
    - Updates C:\Windows\System32\GroupPolicy\Machine\Microsoft\Windows NT\Audit\audit.csv
    - Policies will be visible in secpol.msc → Advanced Audit Policy Configuration
    - WARNING: Domain GPO may overwrite these settings on domain-joined systems

#>

param(
    [Parameter(Position=0, ValueFromRemainingArguments=$true)]
    [string[]]$PolicyValuesArray = @("[]"),
    
    [ValidateSet('Silent','Normal','Verbose','Debug')]
    [string]$LogLevel = 'Normal',
    
    [string]$LogPath = $null,
    
    [switch]$WhatIf
)

# Combine all arguments into a single PolicyValues string
# First, try to get the original command line with proper quotes
$PolicyValues = $null
try {
    $currentPID = $PID
    Write-Host "Current Process ID: $currentPID" -ForegroundColor Cyan
    
    $process = Get-CimInstance Win32_Process -Filter "ProcessId = $currentPID"
    
    if ($process) {
        $commandLine = $process.CommandLine
        Write-Host "Full command line: $commandLine" -ForegroundColor Yellow
        
        if ($commandLine) {
            # Get the script name for more precise regex matching
            $scriptName = [System.IO.Path]::GetFileName($MyInvocation.MyCommand.Path)
            $escapedScriptName = [regex]::Escape($scriptName)
            
            # Extract the first argument after this specific script (with all quotes intact)
            # Stop at known parameters: -LogLevel, -LogPath, -WhatIf, or end of string
            $pattern = "-File\s+`"[^`"]*\\$escapedScriptName`"\s+(.+?)(?:\s+(?:-LogLevel|-LogPath|-WhatIf)|$)"
            Write-Host "Using regex pattern: $pattern" -ForegroundColor DarkGray
            
            if ($commandLine -match $pattern) {
                $rawArgument = $matches[1].Trim()
                Write-Host "Raw argument extracted: $rawArgument" -ForegroundColor Magenta
                
                # Remove outer quotes if present
                if ($rawArgument -match '^"(.*)"$') {
                    $PolicyValues = $matches[1]
                } else {
                    $PolicyValues = $rawArgument
                }
                Write-Host "Extracted PolicyValues from command line: $PolicyValues" -ForegroundColor Green
            } else {
                Write-Host "Command line regex did not match. Command line: $commandLine" -ForegroundColor Red
            }
        } else {
            Write-Host "CommandLine property is null or empty" -ForegroundColor Red
        }
    } else {
        Write-Host "Failed to get process information for PID $currentPID" -ForegroundColor Red
    }
} catch {
    Write-Host "Error extracting from command line: $($_.Exception.Message)" -ForegroundColor Red
    Write-Verbose "Could not extract from command line: $($_.Exception.Message)"
}

# Fallback: Use parameter-based approach if command line extraction failed
if (-not $PolicyValues) {
    $PolicyValues = if ($PolicyValuesArray.Count -gt 1) {
        # Multiple arguments - join them back together
        $PolicyValuesArray -join ''
    } else {
        # Single argument - use as-is
        $PolicyValuesArray[0]
    }
    Write-Verbose "Using parameter-based PolicyValues: $PolicyValues"
}

# Logon/Logoff Policy Database (corrected - removed duplicate)
$PolicyDatabase = @(
    @{ Name = "Audit Logon";                        KeyGroup = "Audit"; Key = "Logon" },
    @{ Name = "Audit Logoff";                       KeyGroup = "Audit"; Key = "Logoff" },
    @{ Name = "Audit Account Lockout";              KeyGroup = "Audit"; Key = "Account Lockout" },
    @{ Name = "Audit IPsec Main Mode";              KeyGroup = "Audit"; Key = "IPsec Main Mode" },
    @{ Name = "Audit IPsec Quick Mode";             KeyGroup = "Audit"; Key = "IPsec Quick Mode" },
    @{ Name = "Audit IPsec Extended Mode";          KeyGroup = "Audit"; Key = "IPsec Extended Mode" },
    @{ Name = "Audit Special Logon";                KeyGroup = "Audit"; Key = "Special Logon" },
    @{ Name = "Audit Other Logon/Logoff Events";    KeyGroup = "Audit"; Key = "Other Logon/Logoff Events" },
    @{ Name = "Audit Network Policy Server";        KeyGroup = "Audit"; Key = "Network Policy Server" },
    @{ Name = "Audit User / Device Claims";         KeyGroup = "Audit"; Key = "User / Device Claims" },
    @{ Name = "Audit Group Membership";             KeyGroup = "Audit"; Key = "Group Membership" },
    @{ Name = "Audit Access Rights";                KeyGroup = "Audit"; Key = "Access Rights" }
)

# Script-wide variables
$script:LogFile = $null
$script:StartTime = Get-Date
$script:ProcessedCount = 0
$script:SuccessCount = 0
$script:FailureCount = 0
$script:SkippedCount = 0

# Initialize logging
function Initialize-LogPath {
    if ($LogPath) {
        $logDir = Split-Path $LogPath -Parent
        if ($logDir -and -not (Test-Path $logDir)) {
            New-Item -ItemType Directory -Path $logDir -Force | Out-Null
        }
        return $LogPath
    }
    
    # Try to get agent directory, fallback to script directory
    $baseDir = $PSScriptRoot
    try {
        $registryPath = if ([Environment]::Is64BitOperatingSystem) {
            "HKLM:\SOFTWARE\WOW6432Node\AdventNet\DesktopCentral\DCAgent"
        } else {
            "HKLM:\SOFTWARE\AdventNet\DesktopCentral\DCAgent"
        }
        
        $agentDir = Get-ItemProperty -Path $registryPath -Name "DCAgentInstallDir" -ErrorAction SilentlyContinue | 
                    Select-Object -ExpandProperty DCAgentInstallDir
        
        if ($agentDir -and (Test-Path $agentDir)) {
            $baseDir = $agentDir
        }
    } catch {
        Write-Verbose "Using script directory for logs"
    }
    
    # Create log directory and file path
    $auditDir = Join-Path (Join-Path $baseDir "logs") "SecurityPolicies"
    if (-not (Test-Path $auditDir)) {
        New-Item -ItemType Directory -Path $auditDir -Force | Out-Null
    }
    
    $timestamp = Get-Date -Format "yyyyMMdd_HHmmss"
    $scriptName = [System.IO.Path]::GetFileNameWithoutExtension($MyInvocation.ScriptName)
    if ([string]::IsNullOrEmpty($scriptName)) {
        $scriptName = "Set-LogonLogoffAuditPolicies"
    }
    return Join-Path $auditDir "${scriptName}_$timestamp.log"
}

$script:LogFile = try { Initialize-LogPath } catch { $null }

# Logging Functions
function Write-Log {
    param(
        [Parameter(Mandatory=$true)]
        [string]$Message,
        
        [ValidateSet('Info','Warning','Error','Debug','Success')]
        [string]$Level = 'Info',
        
        [string]$Component = 'LogonLogoffAudit'
    )
    
    $timestamp = Get-Date -Format 'yyyy-MM-dd HH:mm:ss'
    $logMessage = "[$timestamp] [$Level] [$Component] $Message"
    
    # Console output based on level and LogLevel setting
    switch ($Level) {
        'Error' { 
            if ($LogLevel -ne 'Silent') { Write-Error $Message }
        }
        'Warning' { 
            if ($LogLevel -notin @('Silent')) { Write-Warning $Message }
        }
        'Success' {
            if ($LogLevel -notin @('Silent')) { Write-Host $Message -ForegroundColor Green }
        }
        'Debug' {
            if ($LogLevel -eq 'Debug') { Write-Host $Message -ForegroundColor Gray }
        }
        'Info' {
            if ($LogLevel -notin @('Silent')) { Write-Host $Message }
        }
    }
    
    # File output if LogPath is specified
    if ($script:LogFile) {
        Add-Content -Path $script:LogFile -Value $logMessage -Encoding UTF8
    }
}

function Write-ProgressLog {
    param(
        [int]$Current,
        [int]$Total,
        [string]$Activity = "Processing Logon/Logoff Audit Policies",
        [string]$CurrentItem = ""
    )
    
    if ($LogLevel -ne 'Silent') {
        $percentComplete = if ($Total -gt 0) { ($Current / $Total) * 100 } else { 0 }
        Write-Progress -Activity $Activity -Status "Processing $Current of $Total - $CurrentItem" -PercentComplete $percentComplete
    }
}

function Test-Admin {
    $id = [Security.Principal.WindowsIdentity]::GetCurrent()
    $p  = New-Object Security.Principal.WindowsPrincipal($id)
    return $p.IsInRole([Security.Principal.WindowsBuiltinRole]::Administrator)
}

# Update the Group Policy audit database (audit.csv) so settings show in secpol.msc
# Update the Group Policy audit database (audit.csv) so settings show in secpol.msc
function Update-AuditDatabase {
    param(
        [Parameter(Mandatory=$true)]
        [string]$Subcategory,
        
        [Parameter(Mandatory=$true)]
        [string]$PolicyName,
        
        [Parameter(Mandatory=$false)]
        [ValidateSet('enable', 'disable')]
        [string]$Success,
        
        [Parameter(Mandatory=$false)]
        [ValidateSet('enable', 'disable')]
        [string]$Failure
    )
    
    $auditDbPath = "$env:SystemRoot\System32\GroupPolicy\Machine\Microsoft\Windows NT\Audit"
    $auditCsvPath = Join-Path $auditDbPath "audit.csv"
    
    # Create directory if it doesn't exist
    if (-not (Test-Path $auditDbPath)) {
        Write-Log "Creating audit database directory: $auditDbPath" -Level Info -Component "AuditDB"
        if (-not $WhatIf) {
            New-Item -Path $auditDbPath -ItemType Directory -Force | Out-Null
        }
    }
    
    # Get GUID from auditpol for this subcategory
    $subcategoryGuid = ""
    try {
        $auditpolOutput = & auditpol /list /subcategory:* /r 2>$null
        if ($auditpolOutput) {
            # Parse CSV output and find matching subcategory
            $guidLine = $auditpolOutput | Select-String -Pattern "^\s*$Subcategory," 
            if ($guidLine) {
                # Extract GUID from the line (format: "Subcategory Name,{GUID}")
                if ($guidLine.Line -match '\{([0-9A-F-]+)\}') {
                    $subcategoryGuid = $matches[0]  # Include braces
                    Write-Log "Retrieved GUID for '$Subcategory': $subcategoryGuid" -Level Debug -Component "AuditDB"
                }
            }
        }
    } catch {
        Write-Log "Could not retrieve GUID for '$Subcategory': $($_.Exception.Message)" -Level Debug -Component "AuditDB"
    }
    
    # Initialize or read existing audit.csv
    $csvHeader = "Machine Name,Policy Target,Subcategory,Subcategory GUID,Inclusion Setting,Exclusion Setting,Setting Value"
    $auditEntries = [System.Collections.ArrayList]@()
    $hasHeader = $false
    $foundEntry = $false  # Track if we found entry for this subcategory
    $existingGuid = $subcategoryGuid  # Use the GUID we just retrieved
    $existingSubcategoryName = ""  # Store the subcategory name with prefix
    
    if (Test-Path $auditCsvPath) {
        Write-Log "Reading existing audit database: $auditCsvPath" -Level Debug -Component "AuditDB"
        # Read with UTF8 encoding
        $csvLines = [System.IO.File]::ReadAllLines($auditCsvPath, [System.Text.UTF8Encoding]::new($true))
        
        foreach ($line in $csvLines) {
            # Skip header line
            if ($line -match '^Machine Name,') {
                $hasHeader = $true
                continue
            }
            
            # Skip empty lines
            if ([string]::IsNullOrWhiteSpace($line)) {
                continue
            }
            
            # Check if this line is a policy entry (starts with comma for Machine Name field)
            if ($line -match '^,([^,]*),([^,]+),([^,]*)') {
                $lineSubcategory = $matches[2].Trim()
                $lineGuid = $matches[3].Trim()
                
                # Check if this is our subcategory (with or without "Audit" prefix)
                $isOurSubcategory = ($lineSubcategory -eq $Subcategory) -or ($lineSubcategory -eq "Audit $Subcategory")
                
                if ($isOurSubcategory) {
                    # Found our subcategory - save GUID if it exists and skip (we'll add updated version)
                    $foundEntry = $true
                    if (-not [string]::IsNullOrEmpty($lineGuid)) {
                        $existingGuid = $lineGuid
                    }
                    $existingSubcategoryName = $lineSubcategory
                    Write-Log "Found existing entry for: $lineSubcategory (GUID: $lineGuid) - will update it" -Level Debug -Component "AuditDB"
                    continue
                } else {
                    # Preserve other policy entries
                    Write-Log "Preserving existing entry for: $lineSubcategory" -Level Debug -Component "AuditDB"
                    $auditEntries.Add($line) | Out-Null
                }
            } else {
                # Preserve non-policy lines (like Option: entries)
                $auditEntries.Add($line) | Out-Null
            }
        }
    } else {
        Write-Log "Audit database does not exist, will create: $auditCsvPath" -Level Info -Component "AuditDB"
    }
    
    # Build audit setting value (0=None, 1=Success, 2=Failure, 3=Both)
    # and inclusion setting text
    $auditValue = 0
    $inclusionSetting = "No Auditing"
    
    if ($Success -eq 'enable' -and $Failure -eq 'enable') {
        $auditValue = 3  # Both
        $inclusionSetting = "Success and Failure"
        $valueDesc = "Success and Failure"
    } elseif ($Success -eq 'enable') {
        $auditValue = 1  # Success only
        $inclusionSetting = "Success"
        $valueDesc = "Success"
    } elseif ($Failure -eq 'enable') {
        $auditValue = 2  # Failure only
        $inclusionSetting = "Failure"
        $valueDesc = "Failure"
    } else {
        $inclusionSetting = "No Auditing"
        $valueDesc = "No Auditing"
    }
    
    # Determine the subcategory name to use (with "Audit " prefix)
    # If existing entry had a name, use it; otherwise add "Audit " prefix
    if ($foundEntry -and $existingSubcategoryName) {
        $csvSubcategoryName = $existingSubcategoryName
    } else {
        $csvSubcategoryName = "Audit $Subcategory"
    }
    
    # Format: Machine Name,Policy Target,Subcategory,Subcategory GUID,Inclusion Setting,Exclusion Setting,Setting Value
    # Example: ,System,Audit Security State Change,{0cce9210-69ae-11d9-bed3-505054503030},Success,,1
    
    # Create the entry with GUID
    $updatedEntry = ",System,$csvSubcategoryName,$existingGuid,$inclusionSetting,,$auditValue"
    $auditEntries.Add($updatedEntry) | Out-Null
    Write-Log "Added entry for '$csvSubcategoryName' with GUID ${existingGuid}: $valueDesc (Value=$auditValue)" -Level Info -Component "AuditDB"
    
    if (-not $WhatIf) {
        # Write updated audit.csv with header
        $outputLines = @()
        $outputLines += $csvHeader
        $outputLines += $auditEntries
        
        # Use UTF8 with BOM for better compatibility
        $utf8BOM = New-Object System.Text.UTF8Encoding $true
        [System.IO.File]::WriteAllLines($auditCsvPath, $outputLines, $utf8BOM)
        
        Write-Log "Audit database updated successfully" -Level Success -Component "AuditDB"
        Write-Log "Policy '$PolicyName' will now show in secpol.msc UI" -Level Success -Component "AuditDB"
    } else {
        Write-Log "WHATIF: Would update audit database for '$Subcategory' with value $auditValue ($inclusionSetting)" -Level Info -Component "AuditDB"
    }
}

function Initialize-Script {
    Write-Log "========== Logon/Logoff Audit Policy Framework Script Started =========="
    Write-Log "Script: $($MyInvocation.MyCommand.Name)"
    Write-Log "User: $env:USERNAME"
    Write-Log "Computer: $env:COMPUTERNAME"
    Write-Log "PowerShell Version: $($PSVersionTable.PSVersion)"
    Write-Log "Log Level: $LogLevel"
    
    if ($script:LogFile) {
        Write-Log "Log File: $($script:LogFile)"
    } else {
        Write-Log "Logging: Console only"
    }
    
    if ($WhatIf) {
        Write-Log "WhatIf mode enabled - no changes will be applied" -Level Warning
    }
    
    if (-not (Test-Admin)) {
        Write-Log "Administrator privileges required. Please run this script as Administrator." -Level Error
        throw "Administrator privileges required"
    }
    
    Write-Log "Administrator check passed" -Level Success
    
    # Validate PolicyDatabase
    if (-not $PolicyDatabase -or $PolicyDatabase.Count -eq 0) {
        Write-Log "PolicyDatabase is empty. Please configure audit policies before running." -Level Warning
        return $false
    }
    
    # Validate all policies are audit type
    $nonAuditPolicies = $PolicyDatabase | Where-Object { $_.KeyGroup -ne "Audit" }
    if ($nonAuditPolicies) {
        Write-Log "Found $($nonAuditPolicies.Count) non-audit policies in database. This script only handles audit policies." -Level Error
        return $false
    }
    
    # Validate required properties
    foreach ($policy in $PolicyDatabase) {
        $hasName = $policy.ContainsKey('Name') -or ($policy.PSObject.Properties.Name -contains 'Name')
        $hasKeyGroup = $policy.ContainsKey('KeyGroup') -or ($policy.PSObject.Properties.Name -contains 'KeyGroup')
        $hasKey = $policy.ContainsKey('Key') -or ($policy.PSObject.Properties.Name -contains 'Key')
        
        if (-not ($hasName -and $hasKeyGroup -and $hasKey)) {
            $missingProps = @()
            if (-not $hasName) { $missingProps += 'Name' }
            if (-not $hasKeyGroup) { $missingProps += 'KeyGroup' }
            if (-not $hasKey) { $missingProps += 'Key' }
            Write-Log "Policy '$($policy.Name)' missing required properties: $($missingProps -join ', ')" -Level Error
            return $false
        }
    }
    
    Write-Log "PolicyDatabase validation passed - contains $($PolicyDatabase.Count) Logon/Logoff audit policies" -Level Success
    return $true
}

if (-not (Initialize-Script)) {
    return
}

# Audit Policy Functions
function Set-AuditPolicy {
    param(
        [Parameter(Mandatory=$true)]
        [string]$PolicyName,
        
        [Parameter(Mandatory=$true)]
        [string]$Subcategory,
        
        [Parameter(Mandatory=$false)]
        [ValidateSet('enable', 'disable')]
        [string]$Success,
        
        [Parameter(Mandatory=$false)]
        [ValidateSet('enable', 'disable')]
        [string]$Failure
    )
    
    Write-Log "Setting audit policy: $PolicyName (Subcategory: $Subcategory)" -Level Debug -Component "Audit"
    
    if ([string]::IsNullOrWhiteSpace($Subcategory)) {
        throw "Subcategory name is required for audit policy: $PolicyName"
    }

    $auditArgs = @('/set', "/subcategory:`"$Subcategory`"")
    
    if ($Success) { 
        $auditArgs += "/success:$Success"
        Write-Log "Success auditing: $Success" -Level Debug -Component "Audit"
    }
    if ($Failure) { 
        $auditArgs += "/failure:$Failure" 
        Write-Log "Failure auditing: $Failure" -Level Debug -Component "Audit"
    }

    if ($auditArgs.Count -le 2) {
        throw "At least one of Success or Failure auditing must be specified for policy: $PolicyName"
    }

    $display = "auditpol.exe " + ($auditArgs -join ' ')
    Write-Log "Audit command: $display" -Level Debug -Component "Audit"

    if (-not $WhatIf) {
        Write-Log "Executing: $display" -Level Info -Component "Audit"
        $p = Start-Process -FilePath 'auditpol.exe' -ArgumentList $auditArgs -NoNewWindow -Wait -PassThru
        
        if ($p.ExitCode -ne 0) {
            Write-Log "auditpol failed (exit $($p.ExitCode)) for subcategory '$Subcategory'" -Level Error -Component "Audit"
            throw "auditpol failed (exit $($p.ExitCode)) for subcategory '$Subcategory'"
        } else {
            $auditSettings = @()
            if ($Success) { $auditSettings += "Success:$Success" }
            if ($Failure) { $auditSettings += "Failure:$Failure" }
            $settingsText = $auditSettings -join ", "
            Write-Log "Successfully applied audit policy for '$Subcategory' ($settingsText)" -Level Success -Component "Audit"
        }
    } else {
        Write-Log "WHATIF: Would execute: $display" -Level Info -Component "Audit"
    }
}

function Invoke-AuditPolicy {
    param(
        [Parameter(Mandatory=$true)]
        [pscustomobject]$Policy,
        
        [AllowEmptyString()]
        [string]$Argument = ""
    )
    
    Write-Log "Processing audit policy: '$($Policy.Name)' with argument: '$Argument'" -Level Debug -Component "Audit"
    
    # Handle simplified input format
    # Empty/null argument means disable both Success and Failure auditing
    
    $successAudit = $null
    $failureAudit = $null
    
    if ([string]::IsNullOrEmpty($Argument)) {
        # Disable both success and failure auditing
        $successAudit = 'disable'
        $failureAudit = 'disable'
        Write-Log "Configuring NO AUDITING for '$($Policy.Name)' (disabling both Success and Failure)" -Level Info -Component "Audit"
    } else {
        switch ($Argument) {
            '1' {
                # Success only - explicitly disable failure auditing
                $successAudit = 'enable'
                $failureAudit = 'disable'
                Write-Log "Configuring SUCCESS auditing for '$($Policy.Name)' (disabling Failure)" -Level Info -Component "Audit"
            }
            '0' {
                # Failure only - explicitly disable success auditing
                $successAudit = 'disable'
                $failureAudit = 'enable'
                Write-Log "Configuring FAILURE auditing for '$($Policy.Name)' (disabling Success)" -Level Info -Component "Audit"
            }
            '1,0' {
                # Both success and failure
                $successAudit = 'enable'
                $failureAudit = 'enable'
                Write-Log "Configuring SUCCESS AND FAILURE auditing for '$($Policy.Name)'" -Level Info -Component "Audit"
            }
            '0,1' {
                # Both success and failure
                $successAudit = 'enable'
                $failureAudit = 'enable'
                Write-Log "Configuring SUCCESS AND FAILURE auditing for '$($Policy.Name)'" -Level Info -Component "Audit"
            }
            default {
                Write-Log "Invalid audit format for '$($Policy.Name)': '$Argument'. Expected '1', '0', '1,0', or ''. Skipping..." -Level Warning -Component "Audit"
                return
            }
        }
    }

    # Step 1: Set active audit policy via auditpol
    $setAuditParams = @{
        PolicyName = $Policy.Name
        Subcategory = $Policy.Key
    }
    
    if ($successAudit) { $setAuditParams.Success = $successAudit }
    if ($failureAudit) { $setAuditParams.Failure = $failureAudit }
    
    Set-AuditPolicy @setAuditParams
    
    # Step 2: Update audit database so policy shows in secpol.msc UI
    $updateDbParams = @{
        Subcategory = $Policy.Key
        PolicyName = $Policy.Name
    }
    
    if ($successAudit) { $updateDbParams.Success = $successAudit }
    if ($failureAudit) { $updateDbParams.Failure = $failureAudit }
    
    Update-AuditDatabase @updateDbParams
}

# Main processing loop with comprehensive logging
Write-Log "Starting Logon/Logoff audit policy processing for $($PolicyDatabase.Count) policies"
# Parse PolicyValues array string to array
function Parse-PolicyValuesArray {
    param([string]$ArrayString)
    
    # Check if input is JSON format (starts with [ and ends with ])
    if ($ArrayString -match '^\s*\[.*\]\s*$') {
        Write-Log "Detected JSON format input, attempting to parse..." -Level Debug
        try {
            # Parse the JSON array string
            $Arguments = ConvertFrom-Json $ArrayString
            Write-Log "Successfully parsed JSON policy values array: $($Arguments.Count) values provided" -Level Info
            return $Arguments
        }
        catch {
            Write-Log "Failed to parse as JSON: $($_.Exception.Message)" -Level Warning
            Write-Log "Falling back to positional argument parsing..." -Level Info
        }
    } else {
        Write-Log "Input is not in JSON format (doesn't start with [ ), using as positional arguments" -Level Info
    }
    
    # Fallback: Use PolicyValuesArray as positional arguments
    Write-Log "Using $($PolicyValuesArray.Count) positional arguments" -Level Info
    return $PolicyValuesArray
}
Write-Log "Arguments provided: $PolicyValues"
# Parse PolicyValues array string to get individual arguments
$Arguments = Parse-PolicyValuesArray -ArrayString $PolicyValues

Write-Log "Arguments provided: $($Arguments.Count)"

for ($i = 0; $i -lt $PolicyDatabase.Count; $i++) {
    $script:ProcessedCount++
    $policy = $PolicyDatabase[$i]
    
    Write-ProgressLog -Current ($i + 1) -Total $PolicyDatabase.Count -CurrentItem $policy.Name
    
    # Use empty string if no argument provided (will disable both Success and Failure)
    $arg = if ($i -ge $Arguments.Count) { "" } else { ([string]$Arguments[$i]).Trim() }
   
    Write-Log "Processing policy #$($i+1): '$($policy.Name)' with argument: '$arg'" -Level Info
    
    # Show policy details
    if ($LogLevel -eq 'Debug') {
        Write-Log "Policy Details - Name: $($policy.Name), KeyGroup: $($policy.KeyGroup), Key: $($policy.Key), Argument: $arg" -Level Debug
    }

    try {
        Invoke-AuditPolicy -Policy $policy -Argument $arg
        $script:SuccessCount++
        Write-Log "Successfully processed audit policy: '$($policy.Name)'" -Level Success
    }
    catch {
        $script:FailureCount++
        Write-Log "Failed to process audit policy '$($policy.Name)': $($_.Exception.Message)" -Level Error
        
        # Continue processing other policies unless it's a critical error
        if ($_.Exception.Message -match "Administrator|Permission|Access|auditpol") {
            Write-Log "Critical audit error encountered. Check permissions and auditpol availability." -Level Error
            # Don't break - continue with other policies
        }
    }
    
    Write-Log "--- Policy #$($i+1) completed ---" -Level Debug
}

# Final summary and cleanup
function Write-CompletionSummary {
    $endTime = Get-Date
    $duration = $endTime - $script:StartTime
    
    Write-Log "========== Logon/Logoff Audit Policy Framework Execution Summary ==========" -Level Info
    Write-Log "Execution Duration: $($duration.ToString('hh\:mm\:ss'))" -Level Info
    Write-Log "Total Policies in Database: $($PolicyDatabase.Count)" -Level Info
    Write-Log "Successfully Applied: $script:SuccessCount" -Level Success
    Write-Log "Failed: $script:FailureCount" -Level $(if ($script:FailureCount -gt 0) { 'Warning' } else { 'Info' })
    Write-Log "Not Configured: $script:SkippedCount" -Level Info
    
    $actuallyProcessed = $script:SuccessCount + $script:FailureCount
    if ($actuallyProcessed -gt 0) {
        $successRate = [math]::Round(($script:SuccessCount / $actuallyProcessed) * 100, 2)
        Write-Log "Success Rate: $successRate% (of actually processed policies)" -Level Info
    }
    
    if ($script:FailureCount -gt 0) {
        Write-Log "Some Logon/Logoff audit policies failed to apply. Check the log for details." -Level Warning
        Write-Log "Common issues: Invalid subcategory names, insufficient permissions, auditpol not available" -Level Warning
    }
    
    if ($script:SuccessCount -gt 0) {
        Write-Log " " -Level Info
        Write-Log "VERIFICATION STEPS:" -Level Info
        Write-Log "1. Run: auditpol /get /category:'Logon/Logoff'" -Level Info
        Write-Log "2. Open: secpol.msc → Advanced Audit Policy Configuration → System Audit Policies → Logon/Logoff" -Level Info
        Write-Log "3. Policies should be visible in both locations" -Level Info
    }
    
    if ($WhatIf) {
        Write-Log "WhatIf mode was enabled - no actual changes were made." -Level Info
    }
    
    if ($script:LogFile) {
        Write-Log "Detailed log saved to: $script:LogFile" -Level Info
    }
    
    Write-Log "========== Logon/Logoff Audit Policy Framework Execution Complete ==========" -Level Info
}

Write-CompletionSummary

# Return exit code based on results
if ($script:FailureCount -gt 0) {
    exit 1
} else {
    exit 0
}